clc;clear;close all;

warning('off')
% filed limits [x_min x_max y_min y_max z_min z_max]
field_limits = [0 100 0 100 0 20];
if isfile("robot.mat")
    load("robot.mat");
    [ r_num,state_dim ] = size(robot);            % robot number,state dimensions for each robot
else
    r_num = 4;                                                  % robot number
    state_dim = 3;                                            % state dimensions for each robot
    robot = [ rand(r_num,1)*field_limits(2), rand(r_num,1)*field_limits(4), rand(r_num,1)*field_limits(6) ];
    save robot.mat robot;
end

robot_scatter_size = 10;

f = figure('Name','SimultanuouslyCoverage3D'); hold on;
sc1 = scatter3(robot(:,1),robot(:,2),robot(:,3),robot_scatter_size,'filled'); grid on;
sc3 = scatter3(robot(:,1),robot(:,2),robot(:,3),robot_scatter_size,'filled');
axis(field_limits);view(3);

% discretizing the x-y plane, 
discretization_num = 60; mesh_size = field_limits(2)/discretization_num;
x_space = linspace(field_limits(1),field_limits(2),discretization_num);
y_space = linspace(field_limits(3),field_limits(4),discretization_num);
[X_points,Y_points] = meshgrid(x_space,y_space);
points_xy = [reshape(X_points,discretization_num^2,1) reshape(Y_points,discretization_num^2,1)];

% density calculation
centers = [50 50; 80 80; 20 80; 20 20]; c_index = 1; duration = 1200;
d_center = centers(c_index,:);
R = densityFunc(d_center,points_xy);
%integral = sum(R*mesh_size^2,'all')
h = robot(:,end); sensor_para = 20;
[r,dr_dh] = coverageModel(h,sensor_para,'reciprocal');
%

% initial cost
cost = computeObj_vecform(robot,points_xy,R,mesh_size,r)

lamda_max = 2; lamda_min=0.01; lamda = @(k) (-(lamda_max-lamda_min)/duration *k +lamda_max);
robot_init = robot;
robot_his = cell(100,1);
robot_his{1} = robot;

for k = 1:duration
    robot1 = robot;

    for i = 1:r_num
        % refine the robot set of multual communication
        s_i = robot(i,:); 
        rel_dis = vecnorm(s_i(1:2) - robot(:,1:2),2,2); 
        r = coverageModel(robot(:,end),sensor_para,'reciprocal');
        robot_i = robot; r1 = rel_dis < r + r(i); r1(rel_dis == 0) = 1;
        robot_i = robot_i .* r1;
        [r,dr_dh] = coverageModel(robot_i(:,end),sensor_para,'reciprocal'); r = r.*r1; dr_dh = dr_dh.*r1;
        % compute the gradient
        delta_u_i = computeGradient(i,robot_i,r,dr_dh,points_xy,mesh_size,R,d_center);
        direction = delta_u_i/norm(delta_u_i,2);      
        norm(delta_u_i)
        robot1(i,:) = robot(i,:) + 0.2* direction';
    end
    pause(0.01)
    robot_his{k+1} = robot1;
    robot = robot1;
    h = robot(:,end); [r,dr_dh]= coverageModel(h,sensor_para,'reciprocal');    
    set(sc3,'XData', robot(:,1), 'YData',robot(:,2),'ZData',abs(robot(:,3)));
%     figure(f);scatter3(robot(:,1),robot(:,2),abs(robot(:,3)),robot_scatter_size,'filled');
    drawnow;
    cost = computeObj_vecform(robot,points_xy,R,mesh_size,r);
end


%%
h = robot_init(:,end); 
r = coverageModel(h,20,'reciprocal');
[X,Y,Z] =  computeMeshforPlot(robot_init,field_limits,200);
h1 = figure('Name','CoverageAreaBeforeOptimization'); hold on; grid on;
surf(X,Y,Z); shading interp
axis(field_limits(1:4));

% plotHandles_robot_gt = zeros(r_num,1);
% for i = 1:r_num  
%     x=robot(i,1)-r(i);
%     y=robot(i,2)-r(i);
%     plotHandles_robot_gt(i) = rectangle('Position',[x,y,2*r(i),2*r(i)],'Curvature',[1,1],'LineWidth',1);
% end

%%%%%%%%%%%%%%%%%%%%%       after optimization
robot = reshape(robot,numel(robot)/3,3);
[X,Y,Z] =  computeMeshforPlot(robot,field_limits,200);
h2 = figure('Name','CoverageAreaAfterOptimization'); hold on; grid on;
surf(X,Y,Z); shading interp;colormap default; colorbar
%contour(X,Y,Z,'ShowText','on')
axis(field_limits(1:4));
% plot the coverage area as a circle on the x-y plane
% h = robot(:,end); 
% r = coverageModel(h,20,'reciprocal');
% plotHandles_robot_gt = zeros(r_num,1);
% for i = 1:r_num  
%     x=robot(i,1)-r(i);
%     y=robot(i,2)-r(i);
%     plotHandles_robot_gt(i) = rectangle('Position',[x,y,2*r(i),2*r(i)],'Curvature',[1,1],'LineWidth',1);
% end
%figure(); surf(X_intep,Y_intep,Z_intep); colormap jet
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

